#include <stdio.h>
#include <stdlib.h>
#include <assert.h>

#include "body.h"
#include "cutil_math.h"

#define BLOCK_SIZE 8

#define G 6.67384e-11f

__global__ void updateVel(float4* pos, float3* vel, int n, float dt)
{
	int x = blockIdx.x*BLOCK_SIZE + threadIdx.x;

	if (x < n)
	{
		float4 pos_a = pos[x];
		float3 f = make_float3(0.0f);

		// sum total force
		int i;
		for (i = 0; i != n; ++i)
		{
			if (i != x)
			{
				float4 pos_b = pos[i];
				float3 r = make_float3(pos_a) - make_float3(pos_b);	// 3 flops
				float dis_2 = dot(r, r);	// 5 flops
				float inv_dist_3 = rsqrtf(dis_2*dis_2*dis_2); // 3 flops
				f += (-G*(pos_a.w*pos_b.w)*inv_dist_3)*r; // 6 flops
			}
		}

		// apply acceleration
		vel[x] += (dt/pos_a.w)*f;	// 7 flops
	}
}

__global__ void updatePos(float4* pos, float3* vel, int n, float dt)
{
	int x = blockIdx.x*BLOCK_SIZE + threadIdx.x;

	if (x < n)
	{
		pos[x] += make_float4(dt*vel[x]);	// 3 flops
	}
}

int main()
{
	bodyList* bl = readBodyList();

	//dim3 dimb(BLOCK_SIZE, BLOCK_SIZE);
	dim3 dimb(BLOCK_SIZE, 1);
	dim3 dimg((bl->size-1)/dimb.x + 1, 1);

	float4* d_pos;
	float3* d_vel;

	cudaMalloc(&d_pos, bl->size*sizeof(float4));
	cudaMalloc(&d_vel, bl->size*sizeof(float3));
	cudaMemcpy(d_pos, bl->pos, bl->size*sizeof(float4), cudaMemcpyHostToDevice);
	cudaMemcpy(d_vel, bl->vel, bl->size*sizeof(float3), cudaMemcpyHostToDevice);

	updateVel<<<dimg, dimb>>>(d_pos, d_vel, bl->size, 1.0f);
	cudaThreadSynchronize();
	updatePos<<<dimg, dimb>>>(d_pos, d_vel, bl->size, 1.0f);

	cudaMemcpy(bl->pos, d_pos, bl->size*sizeof(float4), cudaMemcpyDeviceToHost);
	cudaMemcpy(bl->vel, d_vel, bl->size*sizeof(float3), cudaMemcpyDeviceToHost);

	// print output
	int i;
	for (i = 0; i != bl->size; ++i)
	{
		float4* pos = &bl->pos[i];
		float3* vel = &bl->vel[i];
		printf("%9.2f %9.2f %9.2f %9.2f %9.2f %9.2f %9.2f\n", pos->x, pos->y, pos->z,
		//printf("%9.2e %9.2e %9.2e %9.2e %9.2e %9.2e %9.2e\n", pos->x, pos->y, pos->z,
			vel->x, vel->y, vel->z, pos->w);
	}

	cudaFree(d_pos);
	cudaFree(d_vel);

	freeBodyList(bl);
}
